import numpy as np
import matplotlib.pyplot as plt
import h5py
import argparse

pars = argparse.ArgumentParser()
pars.add_argument('-model', type=str, default='A2', help='模型名称')
args = pars.parse_args()

def process_and_plot_ds_file(filename='A2.ds', model_name='A2'):
    """
    处理.ds文件，绘制图像并保存为HDF5格式
    考虑MATLAB和NumPy之间矩阵存储顺序的差异
    
    Parameters:
    filename: .ds文件名
    model_name: 模型名称
    """
    
    # 设置参数（与生成grt.dat时相同）
    nf = 800
    nv = 600
    fmin = 0.01
    fmax = 1
    vmin = 5
    vmax = 9
    
    # 创建频率和速度数组
    f = np.linspace(fmin, fmax, nf)
    v = np.linspace(vmin, vmax, nv)
    
    # 读取.ds文件数据
    tmp = np.loadtxt(filename)
    
    data = tmp.T
    # 注意：MATLAB是列优先存储，Python是行优先存储
    # MATLAB中reshape(tmp,nv,nf)等价于Python中reshape(tmp,(nv,nf),order='F')
    data = np.reshape(tmp, (nv, nf), order='F')
    # 绘制图像
    plt.figure(figsize=(10, 6))
    plt.imshow(data, extent=[fmin, fmax, vmin, vmax], aspect='auto', cmap='jet', origin='lower')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Phase Velocity (km/s)')
    plt.title(f'{model_name} Theoretical Dispersion Spectrum')
    plt.colorbar(label='Amplitude')
    plt.xlim(0.01, 1)
    plt.ylim(5, 9)
    plt.savefig(f'{model_name}.png', dpi=300, bbox_inches='tight')
    plt.show()
    

def process_and_plot_npy_file(filename='dispersion_spectrum.npy', model_name='A2'):
    """
    处理.ds文件，绘制图像并保存为HDF5格式
    考虑MATLAB和NumPy之间矩阵存储顺序的差异
    
    Parameters:
    filename: .ds文件名
    model_name: 模型名称
    """
    
    # 设置参数（与生成grt.dat时相同）
    nf = 1000
    nv = 1000
    fmin = 0.01
    fmax = 1
    vmin = 5
    vmax = 9
    
    # 创建频率和速度数组
    f = np.linspace(fmin, fmax, nf)
    v = np.linspace(vmin, vmax, nv)
    
    # 读取.ds文件数据
    tmp = np.abs(np.load(filename))
    
    data = tmp.T
    # 注意：MATLAB是列优先存储，Python是行优先存储
    # MATLAB中reshape(tmp,nv,nf)等价于Python中reshape(tmp,(nv,nf),order='F')
    # data = np.reshape(tmp, (nv, nf), order='F')
    # 绘制图像
    plt.figure(figsize=(10, 6))
    plt.imshow(data, extent=[fmin, fmax, vmin, vmax], aspect='auto', cmap='jet', origin='lower')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Phase Velocity (km/s)')
    plt.title(f'{model_name} Theoretical Dispersion Spectrum')
    plt.colorbar(label='Amplitude')
    plt.xlim(0.01, 1)
    plt.ylim(5, 9)
    plt.savefig(f'{model_name}.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    

# 使用示例
if __name__ == "__main__":
    if args.model=='txt':
        process_and_plot_ds_file(model_name=args.model)

    if args.model=='npy':
        process_and_plot_npy_file(model_name=args.model)
